# 加载一个 reflow 数据集
# 按照 prompt 进行索引，每个prompt 索引到一个 list , list 中每个 ele 是:  (noise,latent), 其在原队列中的idx, 
# 每个 prompt 聚合后的所有 noise 和 latent 计算 ot map , 得到新的编号对应 ; 

# 比如 dataset 中有 6 个 data , 2 prompt, 分别聚合 1,4,5 , 0,2,3
# ot map 之后编号对应变成了 (1,4,5) -> (5,1,4) , (0,2,3) -> (3,2,0)
# 那么总的编号对应变成了 noise(0,1,2,3,4,5) -> latent(3,5,2,0,1,4)
# 只需要存储 (3,5,2,0,1,4) 这个队列，加载每个 latent 的时候需要做一次额外映射

from reflow.data.utils import LMDB_ndarray, data2lmdb
from reflow.data.dataset import get_reflow_dataset, DataPairsWithText
from glob import glob
from tqdm.auto import tqdm
from collections import defaultdict
import numpy as np
import ot

reflow_ds_path="data/coco2014_reflow/tmp20x10"
reflow_ds_cap_paths = sorted(glob(f'{reflow_ds_path}/**/*.txt', recursive=True))
reflow_ds_caps = []
for cap_path in reflow_ds_cap_paths:
    reflow_ds_caps.extend(open(cap_path).read().splitlines())
cap_to_noise_ids = defaultdict(list)
for idx, cap in tqdm(enumerate(reflow_ds_caps)):
    cap_to_noise_ids[cap].append(idx)

ds=DataPairsWithText(
    reflow_ds_path, 
)

def func(ds, indices):
    noises, latents = [], []
    for idx in indices:
        data=ds[idx]
        noises.append(data['noise'])
        latents.append(data['latent'])
    noises = np.stack(noises, axis=0)
    latents = np.stack(latents, axis=0)
    return noises, latents

idx_pairs = []
for prompt in cap_to_noise_ids.keys():
    # 取得每个 prompt 对应的所有 noise(和latent) 的 idx
    noise_indices = cap_to_noise_ids[prompt]
    # 取得所有的 noise 和 latent 并分成两组
    noises, latents = func(ds, noise_indices)
    bs=noises.shape[0]
    noises, latents = noises.reshape(bs,-1), latents.reshape(bs,-1)
    prob_dist_noise, prob_dist_latent=np.ones((bs,))/bs, np.ones((bs,))/bs
    # 计算最优传输
    dist_mat = ot.dist(noises, latents, metric='euclidean')
    G_mat = ot.emd(prob_dist_noise, prob_dist_latent, dist_mat)
    # get 新的 noise -> latent 的 index 对应关系
    noise_indices = np.array(noise_indices)
    latent_indices = np.argmax(G_mat, axis=0)
    latent_indices = noise_indices[latent_indices]
    idx_pairs.extend(list(zip(noise_indices.tolist(), latent_indices.tolist()))) 
    
idx_pairs.sort(key=lambda x : x[0])
print(idx_pairs)